#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdbool.h>
#include <unistd.h>
#include <string.h>
#include <math.h>
#include <sys/time.h>

#include <mpi.h>

#define PRINT_OUTPUT 0
#define PHI (24.0f/100.0f)

void calcLine(double **cgrid, double **ngrid, int lineLen, int start)
{
	double *nv,*cv,*cvd,*cvu,*cvr,*cvl;
	int i = start;
	
	if (start == 0 || start == lineLen-1)
		return;
	
	for (int j=1; j<lineLen-1; j++)
	{
		nv  = &ngrid[i  ][j]; // next value
		cv  = &cgrid[i  ][j]; // current value
		cvu = &cgrid[i+1][j]; // cv-down
		cvd = &cgrid[i-1][j]; // cv-up 
		cvr = &cgrid[i  ][j+1]; // cv-right
		cvl = &cgrid[i  ][j-1]; // cv-left
		
		*nv = *cv + PHI * (-4.0f * *cv + *cvd + *cvu + *cvr + *cvl);
	}
}

void printGrid(double **g, int dim)
{
#if 1
	int x = 0; // view with 0-border
#else
	int x = 1; // view without
#endif

	for (int i=x; i<dim-x; i++)
	{
		for (int j=x; j<dim-x; j++)
		{
			printf("%2.1lf\t", g[i][j]);
		}
		printf("\n");
	}
}

void setHeatSource(double **dst, int dim, double heat)
{
	int heat_from, heat_to;
	heat_from	= dim/4;
	heat_to		= 3*dim/4;

	for (int i=0; i<dim; i++)
	{
		if ( (i >= heat_from) && (i <= heat_to) )
			dst[1][i] = heat;
		else
			dst[1][i] = 0.0f;
	}
}

void usage(char* program)
{
	printf("Syntax Error!\n");
	printf("Usage: %s <iterations> <size> <heat>\n", program);
}

int main(int argc, char *argv[])
{
	/* MPI setup */
	MPI_Comm comm;
	MPI_Status status[2];
	MPI_Request request[2];
	char hostname[50];
	int nodes, rank;

	MPI_Init(&argc, &argv);
	MPI_Comm_rank (MPI_COMM_WORLD, &rank);  // Who am I?
	MPI_Comm_size (MPI_COMM_WORLD, &nodes); // How many processes?
	gethostname(hostname, 50);
	
#if PRINT_OUTPUT > 0
	printf("# I'm process %2d out of %2d (%s)\n", rank, nodes, hostname);
#endif
	MPI_Barrier(MPI_COMM_WORLD);
	
	/* parse parameters */
	if (argc != 4)
	{
		usage(argv[0]);
		MPI_Finalize();
		return -1;
	}
	int iterations	= atoi(argv[1]);
	int size		= atoi(argv[2]);
	double heat		= atof(argv[3]);
	
	/* calculate work distribution */
	int gridDim, gridLine;
	gridLine	= size + 2; // +2 wg. Randproblem
	gridLine	= gridLine - (gridLine % nodes); // adjust problem size to nodes
	gridDim		= gridLine*gridLine;
	
	int linesPerNode, linesRest;
	linesPerNode = gridLine / nodes;
	linesRest = gridLine % nodes;
	
	/* print summary */
#if PRINT_OUTPUT > 0
	if (rank == 0)
	{
		printf("# %2d: iterations=%i size=%i\n", rank, iterations, size);
		printf("# %2d: gridLine=%i gridDim=%i\n", rank, gridLine, gridDim);
		printf("# %2d: linesPerNode=%i linesRest=%i \n", rank, linesPerNode, linesRest);
	}
#endif
	
	/* MPI_Cart init */
	int above, below, ndim;
    int dims[1], periods[1];
    
	ndim = 1; above = 0; below=0;
	dims[0] = nodes; 
    periods[0] = 0;
    
	MPI_Cart_create( MPI_COMM_WORLD, ndim, dims, periods, 0, &comm );
	MPI_Cart_shift( comm, 0, 1, &above, &below );
	//printf("# %2d: above=%i below=%i\n", rank, above, below);
	
	/* allocate memory */
	double **swp = NULL;
	double **cgrid = (double **)malloc(gridLine * sizeof(double*)); // current grid
		for (int i=0; i<gridLine; i++)
			cgrid[i] = (double *)malloc(gridLine * sizeof(double));
			
	double **ngrid = (double **)malloc(gridLine * sizeof(double*)); // current grid
		for (int i=0; i<gridLine; i++)
			ngrid[i] = (double *)malloc(gridLine * sizeof(double));		
	
	for (int i=0; i<gridLine; i++)
	{
		for (int j=0; j<gridLine; j++)
		{
			cgrid[i][j] = 0.0f;
			ngrid[i][j] = 0.0f;
		}
	}

	/* set heat source */
	setHeatSource(cgrid, gridLine, heat);
#if PRINT_OUTPUT > 1
	if (rank == 0) printGrid(cgrid, gridLine);
#endif
	
	/* determine blocks */
	int start = rank * linesPerNode;
	int end   = start + linesPerNode-1;
	//printf("# %2d: start=%i end=%i\n", rank, start, end);
	//sleep(1);

	MPI_Barrier(MPI_COMM_WORLD);
	
	/*
		Init Done
	*/
	
	double *sendAbovePtr = &cgrid[start  ][0];
	double *sendBelowPtr = &cgrid[end    ][0];
	double *recvAbovePtr = &cgrid[start-1][0];
	double *recvBelowPtr = &cgrid[end+1  ][0];
	
	double t[3];
	t[0] = MPI_Wtime();
	for (int i=0; i<iterations; i++)
	{
		//MPI_Wait(&request[0], &status[0]);
		//MPI_Wait(&request[1], &status[1]);
#if 1
		/* optmized version */

		if (i>0)
		{
			if (above != -2) MPI_Recv( recvAbovePtr, gridLine, MPI_DOUBLE, above, 0, comm, &status[0] );
			if (below != -2) MPI_Recv( recvBelowPtr, gridLine, MPI_DOUBLE, below, 0, comm, &status[1] );
		}

		/* pre-calc FIRST line of block */
		calcLine(cgrid, ngrid, gridLine, start+1);
		
		/* pre-calc LAST line of block */
		calcLine(cgrid, ngrid, gridLine, end-1);
		
		if (i<iterations-1)
		{
			if (above != -2) MPI_Isend( sendAbovePtr, gridLine, MPI_DOUBLE, above, 0, comm, &request[0] );
			if (below != -2) MPI_Isend( sendBelowPtr, gridLine, MPI_DOUBLE, below, 0, comm, &request[1] );
		}
		
		/* calc the lines between start and end */
		for (int j=start+1; j<=end; j++)
		{
			calcLine(cgrid, ngrid, gridLine, j);
		}
#else
		/* naiv version */

		if (i>0)
		{
			if (above != -2) MPI_Recv( recvAbovePtr, gridLine, MPI_DOUBLE, above, 0, comm, &status[0] );
			if (below != -2) MPI_Recv( recvBelowPtr, gridLine, MPI_DOUBLE, below, 0, comm, &status[1] );
		}

		for (int j=start; j<=end; j++)
		{
			calcLine(cgrid, ngrid, gridLine, j);
			//printf("# %2d: j=%i\n", rank, j);
		}

		if (i<iterations-1)
		{
			if (above != -2) MPI_Isend( sendAbovePtr, gridLine, MPI_DOUBLE, above, 0, comm, &request[0] );
			if (below != -2) MPI_Isend( sendBelowPtr, gridLine, MPI_DOUBLE, below, 0, comm, &request[1] );
		}
#endif

		/* swap current/next grid */
		swp = ngrid;
		ngrid = cgrid;
		cgrid = swp;
		swp = NULL;
		
		/* keep heat on */
		setHeatSource(cgrid, gridLine, heat);
	}
	t[1] = MPI_Wtime();
	
	/* print matrix */
#if PRINT_OUTPUT > 1
	if (rank == 0) printGrid(cgrid, gridLine);
#endif

	/* print runtime */
	//sleep(1);
	t[2] = t[1]-t[0];
#if PRINT_OUTPUT > 0
	printf("# %2d: runtime=%lf\n", rank, t[2]);
#endif
	
	double runtime = t[2];
	double timePerIt = runtime/iterations;
	
	if(rank == 0)
	{
		fflush(stdout);
		printf("# nodes iterations dim runtime timePerIt\n");
		printf("%i %i %i %lf %lf\n", nodes, iterations, size, runtime, timePerIt);
	}

	/* free memory */
	for (int i=0; i<gridLine; i++)
	{
		free(cgrid[i]);
		free(ngrid[i]);
	}
	free(cgrid);
	free(ngrid);
	
	//printf("# %2d: done\n", rank);
	MPI_Finalize();
	return 0;
}

